import com.shujia.utils.SparkTool
import org.apache.spark.ml.evaluation.RegressionEvaluator
import org.apache.spark.ml.recommendation.{ALS, ALSModel}
import org.apache.spark.sql.{DataFrame, Dataset, Row, SparkSession}

object TestALS {
  def main(args: Array[String]): Unit = {

    // 0::2::3::1424380312
    def parseRating(str: String): Rating = {
      val fields: Array[String] = str.split("::")
      assert(fields.length == 4)
      Rating(fields(0).toInt, fields(1).toInt, fields(2).toFloat, fields(3).toLong)
    }

    val spark: SparkSession = SparkTool.getSparkSession("TestALS")

    import spark.implicits._

    val ratings: DataFrame = spark.read.textFile("data/mllib/als/sample_movielens_ratings.txt")
      .map(parseRating)
      .toDF()

    // 将数据切分成训练集和测试集
    val Array(training, test) = ratings.randomSplit(Array(0.8, 0.2))

    // 创建ALS模型
    val als: ALS = new ALS()
      .setMaxIter(5) // 设置最大的迭代次数
      .setRegParam(0.01) // 设置缩放比例
      .setUserCol("userId") // 用户id列
      .setItemCol("movieId") // 物品id列
      .setRatingCol("rating") // 评分列
    // 使用训练集训练模型
    val model: ALSModel = als.fit(training)

    // 使用测试集测试模型
    val predictions: DataFrame = model.transform(test)

    predictions.show()

    // 使用rmse均方根误差来评估模型
    val evaluator: RegressionEvaluator = new RegressionEvaluator()
      .setMetricName("rmse") // 指定指标的名称
      .setLabelCol("rating") // 指定label标签列
      .setPredictionCol("prediction") // 指定ALS预测结果列
    val rmse: Double = evaluator.evaluate(predictions)
    println(s"Root-mean-square error = $rmse")

    // 进行推荐
    // 给每个用户推荐10部电影
    // Generate top 10 movie recommendations for each user
    val userRecs: DataFrame = model.recommendForAllUsers(10)

    // 给每部电影推荐10个用户
    // Generate top 10 user recommendations for each movie
    val movieRecs: DataFrame = model.recommendForAllItems(10)

    userRecs.show(truncate = false)

    movieRecs.show(truncate = false)


    // 选择特定一批用户/特定一批物品 进行推荐
    // Generate top 10 movie recommendations for a specified set of users
    val users: Dataset[Row] = ratings.select(als.getUserCol).distinct().limit(3)
    val userSubsetRecs: DataFrame = model.recommendForUserSubset(users, 10)
    // Generate top 10 user recommendations for a specified set of movies
    val movies: Dataset[Row] = ratings.select(als.getItemCol).distinct().limit(3)
    val movieSubSetRecs: DataFrame = model.recommendForItemSubset(movies, 10)

    userSubsetRecs.show(truncate = false)
    movieSubSetRecs.show(truncate = false)

  }

  case class Rating(userId: Int, movieId: Int, rating: Float, timestamp: Long)

}
